import tensorflow as tf
import math
from tensorflow.examples.tutorials.mnist import input_data as mnist_data
print("Tensorflow version " + tf.__version__)
tf.set_random_seed(0)
import numpy as np
import cv2

mnist = mnist_data.read_data_sets("data", one_hot=True, reshape=False, validation_size=0)
X = tf.placeholder(tf.float32, [None, 28, 28, 1])
Y_= tf.placeholder(tf.float32, [None, 10])
# variable learning rate
lr= tf.placeholder(tf.float32)
# three convolutional layers with their channel counts, and a
# fully connected layer (tha last layer has 10 softmax neurons)
K = 12  # first convolutional layer output depth
L = 24  # second convolutional layer output depth
M = 36  # third convolutional layer
N = 256 # fully connected layer
W1 = tf.Variable(tf.truncated_normal([5, 5, 1, K], stddev=0.1))  # 5x5 patch, 1 input channel, K output channels
B1 = tf.Variable(tf.ones([K]) /10)
W2 = tf.Variable(tf.truncated_normal([5, 5, K, L], stddev=0.1))
B2 = tf.Variable(tf.ones([L]) /10)
W3 = tf.Variable(tf.truncated_normal([3, 3, L, M], stddev=0.1))
B3 = tf.Variable(tf.ones([M]) /10)
W4 = tf.Variable(tf.truncated_normal([7* 7 *M, N], stddev=0.1))
B4 = tf.Variable(tf.ones([N]) /10)
W5 = tf.Variable(tf.truncated_normal([N,      10], stddev=0.1))
B5 = tf.Variable(tf.ones([10])/10)
# output is 28x28
C1 = tf.nn.conv2d(X,  W1, strides=[1, 1, 1, 1], padding='SAME')
Y1 = tf.nn.relu(C1 + B1)
# output is 14x14
C2 = tf.nn.conv2d(Y1, W2, strides=[1, 1, 1, 1], padding='SAME')
P2 = tf.nn.max_pool(C2, ksize = [1,2,2,1], strides = [1,2,2,1], padding='SAME')
Y2 = tf.nn.relu(P2 + B2)
# output is 7 x 7
C3 = tf.nn.conv2d(Y2, W3, strides=[1, 1, 1, 1], padding='SAME')
P3 = tf.nn.max_pool(C3, ksize = [1,2,2,1], strides = [1,2,2,1], padding='SAME') 
Y3 = tf.nn.relu(P3 + B3)
# reshape the output from the third convolution for the fully connected layer
YY = tf.reshape(Y3, shape=[-1, 7 * 7 * M])
Y4 = tf.nn.relu(tf.matmul(YY, W4) + B4)
Ylogits = tf.matmul(Y4, W5) + B5
Y = tf.nn.softmax(Ylogits)
# cross-entropy loss function (= -sum(Y_i * log(Yi)) ), normalised for batches of 100  images
# TensorFlow provides the softmax_cross_entropy_with_logits function to avoid numerical stability
# problems with log(0) which is NaN
cross_entropy = tf.nn.softmax_cross_entropy_with_logits(logits=Ylogits, labels=Y_)
cross_entropy = tf.reduce_mean(cross_entropy)*100
# accuracy of the trained model, between 0 (worst) and 1 (best)
correct_prediction = tf.equal(tf.argmax(Y, 1), tf.argmax(Y_, 1))
accuracy = tf.reduce_mean(tf.cast(correct_prediction, tf.float32))
# training step, the learning rate is a placeholder
train_step = tf.train.AdamOptimizer(lr).minimize(cross_entropy)
# init
init = tf.global_variables_initializer()
# save model
Saver = tf.train.Saver(max_to_keep = 1)  # defaults to saving all variables
ModelSavePath = './storage_conv/'
def train():
	with tf.Session() as sess:
		sess.run(init)
		ckpt = tf.train.get_checkpoint_state(ModelSavePath)
		if ckpt and ckpt.model_checkpoint_path:
			Saver.restore(sess, ckpt.model_checkpoint_path)
			print '----> restore model ok'
		else:
			print '----> no model so far'
		iterNum = 40000
		for i in xrange(iterNum):
			if (i + 1) % 100 == 0:
				Saver.save(sess, ModelSavePath + "model", global_step = (i + 1))
			# training on batches of 100 images with 100 labels
			batch_X, batch_Y = mnist.train.next_batch(100)
			# learning rate decay
			max_learning_rate = 0.003
			min_learning_rate = 0.0001
			decay_speed = 2000.0 # 0.003-0.0001-2000=>0.9826 done in 5000 iterations
			learning_rate = min_learning_rate + (max_learning_rate - min_learning_rate) * math.exp(-i/decay_speed)
			a, c, s = sess.run([accuracy, cross_entropy, train_step], \
			{X: batch_X, Y_: batch_Y, lr: learning_rate})
			if (i + 1) % 100 == 0:
				print iterNum, 	'%6d'%i, \
								'%9.3f(cross entropy) '%c, \
								'%9.3f(training accuracy) '%a, \
								'%9.5f(learning rate) '%learning_rate
		# after training finished, save W and B into npz file
		np.savez(
			ModelSavePath + 'mnistConvModel.npz', \
			W1 = sess.run(W1), \
			B1 = sess.run(B1), \
			W2 = sess.run(W2), \
			B2 = sess.run(B2), \
			W3 = sess.run(W3), \
			B3 = sess.run(B3), \
			W4 = sess.run(W4), \
			B4 = sess.run(B4), \
			W5 = sess.run(W5), \
			B5 = sess.run(B5))

def test():
	batch_X = np.zeros([6, 28, 28, 1])
	batch_X[0] = np.expand_dims(cv2.imread('../z0.jpg', cv2.IMREAD_GRAYSCALE).astype('float32')/ 255.0, -1)
	batch_X[1] = np.expand_dims(cv2.imread('../z1.jpg', cv2.IMREAD_GRAYSCALE).astype('float32')/ 255.0, -1)
	batch_X[2] = np.expand_dims(cv2.imread('../z2.jpg', cv2.IMREAD_GRAYSCALE).astype('float32')/ 255.0, -1)
	batch_X[3] = np.expand_dims(cv2.imread('../z3.jpg', cv2.IMREAD_GRAYSCALE).astype('float32')/ 255.0, -1)
	batch_X[4] = np.expand_dims(cv2.imread('../z4.jpg', cv2.IMREAD_GRAYSCALE).astype('float32')/ 255.0, -1)
	batch_X[5] = np.expand_dims(cv2.imread('../z5.jpg', cv2.IMREAD_GRAYSCALE).astype('float32')/ 255.0, -1)
	trueY = np.array([6, 7, 2, 4, 3, 5])

	with tf.Session() as sess:
		sess.run(init)
		ckpt = tf.train.get_checkpoint_state(ModelSavePath)
		if ckpt and ckpt.model_checkpoint_path:
			Saver.restore(sess, ckpt.model_checkpoint_path)
			print '----> restore model ok'
			WW = sess.run(W5)
			print WW.shape
		else:
			print '----> no model so far'
		oy = sess.run(Y, {X: batch_X})
		for i in range(0, 6):
			for j in range(0, 10):
				if(oy[i, j] > 0.6):
					print i, trueY[i], j
					break

def test1():
	batch_X, batch_Y = mnist.train.next_batch(1)
	img = batch_X[0,:,:,0] * 255.0
	img = img.astype('uint8')
	print img
#	cv2.imshow('win', img)
#	cv2.waitKey(0)
	print batch_Y
	
	with tf.sess() as sess:
		sess.run(init)
		ckpt = tf.train.get_checkpoint_state(ModelSavePath)
		if ckpt and ckpt.model_checkpoint_path:
			Saver.restore(sess, ckpt.model_checkpoint_path)
			print '----> restore model ok'
		else:
			print '----> no model so far'
		oy = sess.run(Y, {X: batch_X, pkeep: 1.0})
		print oy

if __name__ == "__main__":
    #train()
    test()

